# -*- coding: utf-8 -*-
# standard
import os
import sys
import re
import urllib3
import json
import math
from queue import Queue
from queue import Empty
from platform import system
from threading import Condition
from concurrent.futures import Future
from concurrent.futures import as_completed
from urllib.parse import urlparse
import threading
import _thread
from collections import deque

# third
import requests
import paramiko
from tqdm import tqdm

# local
_P_PATH =  os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
if _P_PATH not in sys.path:
    sys.path.append(_P_PATH)
from models.json_files import *



"""
@Title:  scp 拷贝文件的模型
@File: scp_files.py
@Author: walle 2023年10日17日 10时11分32秒
@Version: 1.0.0
@Desc: 
"""


DEFAULT_BROWSER_HEADERS = {  # 默认浏览器头
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:94.0) Gecko/20100101 Firefox/94.0",
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
}
urllib3.disable_warnings()
# gpu的配置
# USER = "root"  # 文件服务器 ssh 账户 
# PW = "Default01!"# 文件服务器 ssh 密码
# HOST = "gpu1.welltop.tech" # 文件服务器 ssh 地址
# DEST_DIR = "/mnt/aiface/data/i2v/work_data/"   # 目标目录，必须/ 结尾
# hk02的配置
USER = "spider"  # hk02 ssh 账户 
PW = "jingying"# hk02 ssh 密码
if system().lower().startswith("win"):
    HOST = "47.243.81.213" # hk02 ssh 地址 外网地址
else:
    HOST = "172.26.190.232" # hk02 ssh 地址 内网地址 m目前只能内网访问
DEST_DIR = "/mnt/projects/LRS3_m4a/"   # 目标目录，必须/ 结尾

PORT = 22  # 文件服务器 ssh 端口
MIN_VIDEO_SIZE = 1024  # 最小文件尺寸，单位K， 也就是已1m的意思 
STANDARD_QUALITY = 720  # 当前指定的视频质量

K = 1024  # 1K大小
M = 1024 ** 2  # 1m大小
G = 1024 ** 3  # 1g大小


class ClipOrm(Base, OrmTools):
    """
    @Title: 视频片段
    @File: scp_file_type.py
    @Author: walle 2023年10日17日 10时11分32秒
    @Version: 1.0.0
    @Desc: 
    """
    __tablename__ = 'clip'
    id = Column(Integer, primary_key=True)
    start_time = Column(String(64), doc="开始时间")
    end_time = Column(String(64), doc="结束时间")
    # 外键
    media_key = Column(String(64), ForeignKey('media_file.key'), doc="文件")
    media = relationship('MediaFileOrm', back_populates='clips')


class FileDownloadSpeed(BaseModel):
    """
    文件下载的速度对象
    """
    seconds: float = Field(..., description="传输耗时")
    size: float = Field(..., description="传输文件的大小，单位M")
    speed: float = Field(..., description="传输速度, 单位每秒k，")

    @validate_arguments(config={"arbitrary_types_allowed": True})
    def __init__(self, seconds: float, size: int):
        speed = round(size / K / seconds, 2)
        size =  round(size / M, 2)
        super().__init__(seconds=seconds, size=size, speed=speed)



class MediaFileOrm(Base, OrmTools):
    """
    @Title: 视频文件
    @File: scp_file_type.py
    @Author: walle 2023年10日17日 10时11分32秒
    @Version: 1.0.0
    @Desc: 
    """
    __tablename__ = 'media_file'
    key = Column(String(64), primary_key=True)
    caption = Column(String(255), default='', nullable=True, doc="视频标题")
    source_site = Column(String(64), nullable=True, default=VideoSourceSite.YOUTUBE.value, doc="视频来源站点")
    page_url = Column(String(255), nullable=True, default="", doc="视频页面地址")
    resource = Column(JSON, nullable=True, default={}, doc="视频解析出的信息")
    save_path = Column(String(512), nullable=True, default="", doc="视频本地保存路径 空字符分割")
    quality = Column(Integer, nullable=True, default=0, doc="视频质量")
    remote_storage = Column(String(256), nullable=True, default="", doc="视频远程存储地址, ")
    error = Column(String(256), nullable=True, default="", doc="错误信息")
    parse_time = Column(DateTime, nullable=True, default=datetime.now, doc="解析时间")
    need_audio = Column(Integer, default=0, nullable=True, doc="是否还需要下载音频文件")
    audio_path = Column(String(512), nullable=True, default="", doc="音频文件本地保存路径 空字符分割")
    audio_size = Column(Integer, nullable=True, default=0, doc="音频文件尺寸")
    download_begin = Column(DateTime, nullable=True, default=None, doc="视频开始下载时间")
    download_end = Column(DateTime, nullable=True, default=None, doc="视频结束下载时间")
    create_time = Column(DateTime, nullable=True, default=datetime.now, doc="创建时间")

    clips = relationship("ClipOrm", back_populates="media", lazy="dynamic")

    def get_video_info(self) -> dict:
        """
        从接口数据中国
        :raises ValueError: _description_
        :raises ValueError: _description_
        :return: _description_
        """
        data = self.resource
        if "video_info" in data:
            result = data['video_info']
        else:
            result = {}
        return result
    
    def get_speed_obj(self) -> Optional[FileDownloadSpeed]:
        """
        返回一个速度对象，为了方便计算速度, 没有下载成功的返回 None
        :raises ValueError: _description_
        :raises ValueError: _description_
        :return: _description_
        """
        if self.download_begin and self.download_end and self.save_path != "":
            seconds = (self.download_end - self.download_begin).total_seconds()
            video_info = self.get_video_info()
            # 转成字典组成额数组，方便排序
            video_info = [{"quality": int(re.search(r"(\d+)", k).group()), "size": v} for k, v in video_info.items()]
            video_info.sort(key=lambda x: x['quality'], reverse=True)
            size = video_info[0]['size']  # 这个值可能为 None
            speed_obj = FileDownloadSpeed(size=size if size else 0, seconds=seconds)
            return speed_obj
        else:
            pass

    @classmethod
    def generate_save_file_path(cls, key: str, suffix: str = ".mp4") -> str:
        """
        生成保存文件路径
        :param index: 文件索引
        :param file_suffix: 文件后缀
        """
        dir_path = os.path.join(DOWNLOAD_DIR, VideoSourceSite.YOUTUBE.value, date.today().strftime(DateTimeFormat.YMD.value))
        # dir_path = os.path.join(_P_PATH, "work_data")
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
        file_name = f"{key}{suffix}"
        full_path = os.path.join(dir_path, file_name)
        return full_path
    
    @classmethod
    @validate_arguments(config={"arbitrary_types_allowed": True})
    def parse_source_site(cls, download_url: constr(to_lower=True, strip_whitespace=True)) -> Optional[VideoSourceSite]:
        """
        解析下载链接返回源站。(youtube or bilibili)
        :param download_url: json文件里面对应的下载链接
        """
        if matched := re.match("^http\w?://([a-z,-,\.]*)/.*$", download_url):
            domain = matched.group(1)  # 第一个匹配的子组，0是全句。也可以使用 groups()返回子组的元组再用0索引取值
            members = VideoSourceSite.__members__
            source_site = None
            for item in members.values():
                if item.value in domain:
                    source_site = item
                    break
            if source_site:
                return source_site
            else:
                ms = f"域名匹配失败，意料外的域名: {domain}"
                raise ValueError(ms)
        else:
            ms = f"无效的视频url: {download_url}"
            raise ValueError(ms)

    @classmethod
    def reduce_resource(self, resource: dict, key: str):
        """
        把api的解析结果化简，避免超出sqlite的json字段8000的上限
        第二个逻辑是选择尽可能高的分辨率，要求的是最低分辨率720p
        :param resource: _description_
        :return: _description_
        """
        video_info = {}
        video_info2 = {}
        video_info3 = {}
        video_info4 = {}
        for x in resource['medias'][0]['formats']:
            video_info[x['quality_note']] = x['video_size']
            video_info2[x['quality']] = (x['video_url'], x.get("audio_url"))
            video_info3[x['video_url']] = x['quality_note']
            video_info4[x['quality']] = x['video_size']
        # 选择质量最高的视频源
        top_quality = max(video_info2.keys())
        top_quality_video_url, top_quality_audio_url = video_info2[top_quality]
        if top_quality_audio_url is None:
            for _video, audio in video_info2.values():
                if audio is not None:
                    top_quality_audio_url = audio  # 当最高质量视频缺少音频文件时，用低一级的质量的的音频文件替代
                    break
        default_quality_url = resource['medias'][0]['resource_url']
        if top_quality_video_url != default_quality_url:
            default_quality = video_info3[default_quality_url]
            size = video_info4[top_quality]
            logger.info(f"视频 {key} 默认质量为 {default_quality} 不是最高质量，改为 {top_quality}P, 文件尺寸为 {round(size / 1048576)} M")
        result = {
            "caption": resource['text'],
            "resource_url": (top_quality_video_url, top_quality_audio_url),
            "video_info": video_info,
            "current_quality": top_quality,
        }
        return result

    @classmethod
    def analysis_media_url(cls, page_url: str, need_audio: int = 0) -> Resp:
        """
        利用第三方接口从页面地址解析出来 resource(视频的下载信息)
        """
        hhm_api = 'https://h.aaaapp.cn/single_post'     # 单个帖子提取接口 (如果主页批量提取使用：https://h.aaaapp.cn/posts)
        user_id = '5918929AD21EDEC547D4E2E2AFFF8471'         # 用户id
        secret_key = '3faf9d37f6de646aa61a2f371d68fa0e'      # 密钥

        params = {
            'userId': user_id,
            'secretKey': secret_key,
            'url': page_url
        }
        r, resp = None, Resp()
        try:
            r = requests.post(hhm_api, json=params, verify=False)
        except Exception:
            resp.error(traceback.format_exc())
        finally:
            if not resp:
                logger.error(f"请求API接口出错: {resp.message}")
            else:
                logger.debug(f"视频 {page_url} 的接口信息已返回")
                response = r.json()
                if response['code'] != 200:  
                    resp.message = f"API解析失败, 错误原因:{response.get('msg')}"
                    ms = f"使用API解析视频{page_url}的信息失败，status = {response['code']}, reason = {resp.message}"
                    logger.error(ms)
                else:
                    key = urlparse(page_url).query.split("=")[-1]
                    # resp.data = response['data']  # 去掉无用的数据，否则可能超过json 字段的 8000 的长度上限
                    resp.data = cls.reduce_resource(resource=response['data'], key=key)
            return resp



class ScpTask(BaseModel):
    num: int = Field(..., description="序号")
    key: str = Field(..., description="视频的key")
    sources: List[str] = Field([], description="待拷贝的文件")
    

class ScpRecordOrm(Base, OrmTools):
    """
    @Title: 拷贝文件的ORM模型
    @File: scp_files.py
    @Author: walle 2023年10日17日 10时11分32秒
    @Version: 1.0.0
    @Desc: 
    一条记录对应一个文件。
    由于一个 VideoDownloadRecordOrm 对象可能包含多个文件，因此 ScpRecordOrm: VideoDownloadRecordOrm = n: 1
    """
    __tablename__ = 'scp_record'
    id = Column(Integer, primary_key=True, autoincrement=True)
    source = Column(String(512), nullable=False, doc="源文件的路径")
    destination = Column(String(512), nullable=False, doc="拷贝的目标文件的路径")
    copy_begin = Column(DateTime, nullable=True, default=None, doc="拷贝开始时间")
    copy_end = Column(DateTime, nullable=True, default=None, doc="拷贝完成时间")

    @classmethod
    def format_remote_file_info(cls, file_lines: List[str], to_k: bool = False):
        """
        格式化远程文件信息
        :param to_k: 是否转化为k的单位，默认是字节单位
        """
        good_files = {}
        bad_files = {}
        for line in file_lines:
            words = [x for x in line.split(" ") if x.strip() != ""]
            if len(words) < 9:
                if line != "":
                    logger.warning(f"无效的文件信息读取行")
                continue
            try:
                size = int(words[4]) / 1024 if to_k else int(words[4])
            except Exception:
                logger.error(f"无效的文件尺寸: {words[4]} words={words}")
                continue
            file_name = words[8]
            if size < MIN_VIDEO_SIZE:
                logger.warning(f"远程服务器上{file_name}文件尺寸只有{size}k， 可能是因为拷贝过程中断导致的文件不完整")
                bad_files[file_name.rsplit(".", 1)[0]] = size
            else:
                good_files[file_name.rsplit(".", 1)[0]] = size
        return good_files, bad_files
    
    @classmethod
    def view_remote_files(cls, to_k: bool = True, host: str = None, port: int = None, username: str = None, password: str = None, dest_dir: str = None) -> dict:
        """
        浏览远程文件
        :param to_k: 是否转化为k的单位，默认是字节单位
        """
        dest_dir = DEST_DIR if dest_dir is None else dest_dir
        host = HOST if host is None else host
        port = PORT if port is None else port
        username = USER if username is None else username
        password = PW if password is None else password
        logger.debug(f"开始检查远程服务器({host})上的视频文件的完整性...")
        command = f'ls -l {dest_dir} | grep "^-"'
        ssh = paramiko.SSHClient()  
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())  
        # 连接到远程Linux系统  
        ssh.connect(hostname=host, port=port, username=username, password=password)    
        # 执行命令并获取输出  
        stdin, stdout, stderr = ssh.exec_command(command)  
        output = stdout.read().decode()  
        # 打印输出结果  
        # print(output)  
        # 关闭SSH连接  
        ssh.close()
        file_lines = output.split("\n")
        good_files, bad_files = cls.format_remote_file_info(file_lines, to_k=to_k)
        logger.debug(f"远程服务器({host})上的视频文件的完整性检查完毕。完整视频 {len(good_files)} 个， 残缺视频 {len(bad_files)} 个")
        return good_files, bad_files
    
    @classmethod
    def get_need_copy_records(cls, dir_path: str = None) -> List[dict]:
        """
        检查需要拷贝的目录。
        依赖于对目录的检查
        :return: _description_
        """
        data = []
        dir_path = DOWNLOAD_DIR if dir_path is None else dir_path
        for file_name in os.listdir(dir_path):
            current_path = os.path.join(dir_path, file_name)
            if os.path.isfile(current_path):
                key = file_name.rsplit(".", 1)[0]
                stat = os.stat(current_path)
                create_time = datetime.fromtimestamp(stat.st_ctime)
                data.append({"key": key, "create_time": create_time, "sources": [current_path]})
            else:
                data.extend(cls.get_need_copy_records(current_path))
        return data
    
    @classmethod
    def copy_to_remote_forever(cls, tasks: Queue):
        """
        批量拷贝文件向服务器
        当前使用的函数
        :param sources: _description_
        """
        ssh, sftp, prev_task_count = None, None, None
        current_remote_folder = None  
        while 1:
            error = ''
            current_task_count = tasks.qsize()
            if prev_task_count is None or prev_task_count != current_task_count:
                logger.debug(f"拷贝任务队列长度: {current_task_count}")
                prev_task_count = current_task_count
            try:
                task: ScpTask = tasks.get(timeout=1)
            except Empty:
                continue
            logger.info(f"第{task.num}个视频({task.key})的{len(task.sources)}个文件开始拷贝")
            try:
                if ssh is None or sftp is None:
                    ssh = paramiko.SSHClient()
                    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                    logger.debug(f"正在与主机({HOST})建立 ssh 连接....")
                    ssh.connect(hostname=HOST, port=PORT, username=USER, password=PW)
                    sftp = paramiko.SFTPClient.from_transport(ssh.get_transport())
                    logger.debug(f"与主机({HOST})的 ssh 连接建立成功")
                
                
                for source in task.sources:
                    file_name = os.path.basename(source)
                    now = datetime.now()
                    folder = now.strftime("%Y%m%d%H")
                    destination_folder = f"{DEST_DIR}{folder}"
                    destination = f"{destination_folder}/{file_name}"
                    if not os.path.exists(source):
                        logger.warning(f"源文件({source})不存在, 跳过")
                        continue
                    # 检查目标目录是否存在
                    try:
                        sftp.lstat(destination_folder)
                        logger.debug(f"目录 {destination_folder} 已存在")
                    except FileNotFoundError:
                        sftp.mkdir(destination_folder)
                        logger.debug(f"目录 {destination_folder} 已创建")
                    copy_begin = datetime.now()
                    copy_result = cls.copy_one(source=source, destination=destination, ssh=ssh, sftp=sftp)
                    if copy_result:  # 拷贝成功的才会写入数据库
                        DB_THREAD_LOCK.acquire()
                        db_session = new_db_session()
                        copy_end = datetime.now()
                        db_session.add(cls(source=source, destination=destination, copy_begin=copy_begin, copy_end=copy_end))                
                        cls.remove_source(source=source)
                        db_session.commit()
                        db_session.close()
                        DB_THREAD_LOCK.release()
            except Exception:
                error = traceback.format_exc()              
            finally:
                if error == '':
                    tasks.task_done()
                    logger.info(f"第{task.num}个视频({task.key})的{len(task.sources)}个文件拷贝成功")
                else:
                     logger.error(error)
                     logger.error(f"第{task.num}个视频({task.key})的{len(task.sources)}个文件拷贝失败")
        sftp.close()
        ssh.close()

    @classmethod
    def copy_one(cls, source: str, destination: str, ssh: paramiko.SSHClient = None, sftp: paramiko.SFTPClient = None) -> Resp:
        """
        拷贝一个文件
        :param source: _description_
        """
        resp = Resp()
        logger.debug(f"开始向远程服务器写入文件 {destination}")     
        if ssh is None:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname=HOST, port=PORT, username=USER, password=PW)
            sftp = paramiko.SFTPClient.from_transport(ssh.get_transport())
            need_close = True
        else:
            need_close = False
        put_result = None
        locale_data = threading.local()  # 本地变量在线程间独立的保存出数据
        locale_data.prev_completed = 0
        try:
            file_name = os.path.basename(source)
            total = os.path.getsize(source)
            with tqdm(total=total, unit_scale=True, colour="cyan", unit_divisor=1024, unit="B") as progress:

                def progress_bar_func(completed: int, total: int, locale_data: _thread._local, **kwargs):
                    # 进度条函数
                    progress.update(completed - locale_data.prev_completed)
                    locale_data.prev_completed = completed

                put_result = sftp.put(source, destination, callback=partial(progress_bar_func, locale_data=locale_data))  # todo： 这里可能抛出 paramiko.ssh_exception.SSHException: Key-exchange timed out waiting for key negotiation
        except Exception as e:
            resp.error(f"上传文件时发生错误： {e} 本地位置： {source} ; 远程位置： {destination}")
            logger.exception(e)
        finally:
            if put_result and hasattr(put_result, 'st_size'):
                st_size = put_result.st_size
                if st_size == 0:
                    ms = f"文件在远程服务器上的写入好像未成功{destination}"
                    resp.message = ms
                    logger.warning(ms)
                else:
                    ms = f"文件在远程服务器上的写入成功{destination}"
                    logger.debug(ms)
            else:
                ms = f"文件在远程服务器上的写入失败{destination}"
                resp.message = ms
                logger.warning(ms)
            if need_close:
                sftp.close()
                ssh.close()
            return resp

    @classmethod
    def remove_source(cls, source: str):
        """
        删除指定的（源）文件
        :param source: _description_
        """
        os.remove(source)
        try:
            logger.debug(f"删除源文件{source}成功")
        except Exception:
            logger.error(traceback.format_exc())
            logger.error(f"删除源文件{source}失败")
        
    @classmethod
    def get_all_downloaded_video_paths(cls, folder_path: str = None) -> List[str]:
        """
        获取所有的已下载的文件的绝对路径，文件名，尺寸等信息
        :param folder_path: _description_
        :return: _description_
        """
        folder_path = DOWNLOAD_DIR if folder_path is None else folder_path
        data = []
        for file_name in os.listdir(folder_path):
            current_path = os.path.join(folder_path, file_name)
            if file_name.lower().endswith(".mp4") and os.path.isfile(current_path):
                stat = os.stat(current_path)
                info = {
                    "file_path": current_path,
                    "key": file_name.rsplit(".", 1)[0],
                    "file_name": file_name,
                    "file_size": stat.st_size,
                }
                data.append(info)
            elif os.path.isdir(current_path):
                data.extend(cls.get_all_downloaded_video_paths(current_path))
            else:
                pass
        return data
    
    @classmethod
    def check_downloaded_folder(cls, delete_file: bool = False):
        """
        检查已下载的文件。看看他们是否是完整的？不完整的视频删除之
        """
        downloads_videos = cls.get_all_downloaded_video_paths()
        keys = [x['key'] for x in downloads_videos]
        db_session = new_db_session()
        size_info_map = {}
        while len(keys) > 0:
            current, keys = keys[:5000], keys[5000:]
            for x in db_session.query(MediaFileOrm).where(MediaFileOrm.key.in_(current)):
                size_info_map[x.key] = x.get_video_info()
        for download_video in downloads_videos:
            key = download_video['key']
            record = size_info_map.get(key)
            if record:
                current_video_quality_note = None
                for quality_note, video_size in record.items():
                    if video_size == download_video['file_size']:
                        current_video_quality_note = quality_note
                        break
                if current_video_quality_note is None:
                    logger.warning(f"当前视频的大小为{download_video['file_size']},和所有已知质量的视频尺寸都不符。删除! {record}")
                    if delete_file:
                        os.remove(download_video['file_path'])
                else:
                    logger.debug(f"视频{download_video['file_name']}的质量为{current_video_quality_note}, 文件尺寸吻合({download_video['file_size']})")
            else:  # 没有对应记录的视频，检查他们是否大于最小视频长度
                if download_video['file_size'] > MIN_VIDEO_SIZE * 1024:  # 合格的视频
                    pass
                else:
                    logger.warning(f"视频{download_video['file_name']}的尺寸只有 {download_video['file_size']}K, 删除")
                    if delete_file:
                        os.remove(download_video['file_path'])

            
    



if __name__ == '__main__':
    pass